// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2023 KubeRelayer */
#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/in.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

// 定义许可证
char LICENSE[] SEC("license") = "GPL";

// 定义数据结构
struct http_event_t {
    __u64 ts;                // 时间戳
    __u32 pid;               // 进程ID
    __u32 saddr;             // 源IP
    __u32 daddr;             // 目标IP
    __u16 sport;             // 源端口
    __u16 dport;             // 目标端口
    __u8 direction;          // 方向：1=请求，2=响应
    char data[4096];         // 数据
};

// 定义事件映射
struct {
    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
    __uint(key_size, sizeof(int));
    __uint(value_size, sizeof(int));
    __uint(max_entries, 1024);
} http_events SEC(".maps");

// 定义服务过滤映射
struct service_key_t {
    __u32 ip;
    __u16 port;
};

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(key_size, sizeof(struct service_key_t));
    __uint(value_size, sizeof(__u8));
    __uint(max_entries, 1024);
} services SEC(".maps");

// 定义会话跟踪映射
struct session_key_t {
    __u32 saddr;
    __u32 daddr;
    __u16 sport;
    __u16 dport;
};

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(key_size, sizeof(struct session_key_t));
    __uint(value_size, sizeof(__u64));
    __uint(max_entries, 4096);
} sessions SEC(".maps");

// 辅助函数：检查是否为HTTP请求或响应
static __always_inline int is_http(const char *data, int data_len) {
    // 检查是否为HTTP请求
    if (data_len >= 4 && 
        ((data[0] == 'G' && data[1] == 'E' && data[2] == 'T' && data[3] == ' ') ||
         (data[0] == 'P' && data[1] == 'O' && data[2] == 'S' && data[3] == 'T' && data[4] == ' ') ||
         (data[0] == 'P' && data[1] == 'U' && data[2] == 'T' && data[3] == ' ') ||
         (data[0] == 'D' && data[1] == 'E' && data[2] == 'L' && data[3] == 'E' && data[4] == 'T' && data[5] == 'E' && data[6] == ' ') ||
         (data[0] == 'H' && data[1] == 'E' && data[2] == 'A' && data[3] == 'D' && data[4] == ' ') ||
         (data[0] == 'O' && data[1] == 'P' && data[2] == 'T' && data[3] == 'I' && data[4] == 'O' && data[5] == 'N' && data[6] == 'S' && data[7] == ' '))) {
        return 1;
    }
    
    // 检查是否为HTTP响应
    if (data_len >= 4 && 
        data[0] == 'H' && data[1] == 'T' && data[2] == 'T' && data[3] == 'P' && data[4] == '/') {
        return 2;
    }
    
    return 0;
}

// 辅助函数：检查是否为监控的服务
static __always_inline int is_monitored_service(struct service_key_t *key) {
    __u8 *value;
    value = bpf_map_lookup_elem(&services, key);
    return value != NULL;
}

// 辅助函数：检查是否为已知会话
static __always_inline int is_known_session(struct session_key_t *key) {
    __u64 *ts;
    ts = bpf_map_lookup_elem(&sessions, key);
    return ts != NULL;
}

// 辅助函数：添加会话
static __always_inline void add_session(struct session_key_t *key) {
    __u64 ts = bpf_ktime_get_ns();
    bpf_map_update_elem(&sessions, key, &ts, BPF_ANY);
}

// 处理TCP数据包
SEC("socket")
int http_filter(struct __sk_buff *skb) {
    // 验证数据包长度
    if (skb->len < (ETH_HLEN + sizeof(struct iphdr) + sizeof(struct tcphdr))) {
        return 0;
    }
    
    // 获取以太网头
    struct ethhdr eth;
    if (bpf_skb_load_bytes(skb, 0, &eth, sizeof(eth)) < 0) {
        return 0;
    }
    
    // 检查是否是IP包
    if (eth.h_proto != bpf_htons(ETH_P_IP)) {
        return 0;
    }
    
    // 获取IP头
    struct iphdr iph;
    if (bpf_skb_load_bytes(skb, ETH_HLEN, &iph, sizeof(iph)) < 0) {
        return 0;
    }
    
    // 检查是否是TCP包
    if (iph.protocol != IPPROTO_TCP) {
        return 0;
    }
    
    // 获取TCP头
    struct tcphdr tcph;
    int ip_header_len = (iph.ihl & 0xF) << 2;
    int tcp_header_offset = ETH_HLEN + ip_header_len;
    if (bpf_skb_load_bytes(skb, tcp_header_offset, &tcph, sizeof(tcph)) < 0) {
        return 0;
    }
    
    // 获取TCP数据偏移
    int tcp_header_len = (tcph.doff & 0xF) << 2;
    int payload_offset = ETH_HLEN + ip_header_len + tcp_header_len;
    int payload_length = skb->len - payload_offset;
    
    // 检查是否有负载
    if (payload_length <= 0) {
        return 0;
    }
    
    // 创建服务键
    struct service_key_t src_key = {
        .ip = iph.saddr,
        .port = tcph.source
    };
    
    struct service_key_t dst_key = {
        .ip = iph.daddr,
        .port = tcph.dest
    };
    
    // 创建会话键
    struct session_key_t session_key = {
        .saddr = iph.saddr,
        .daddr = iph.daddr,
        .sport = tcph.source,
        .dport = tcph.dest
    };
    
    // 检查是否为监控的服务
    int is_src_monitored = is_monitored_service(&src_key);
    int is_dst_monitored = is_monitored_service(&dst_key);
    
    if (!is_src_monitored && !is_dst_monitored) {
        // 检查是否为已知会话
        if (!is_known_session(&session_key)) {
            // 反向会话键
            struct session_key_t reverse_key = {
                .saddr = iph.daddr,
                .daddr = iph.saddr,
                .sport = tcph.dest,
                .dport = tcph.source
            };
            
            // 检查反向会话
            if (!is_known_session(&reverse_key)) {
                return 0;
            }
        }
    }
    
    // 读取负载数据
    char data[4096] = {0};
    int data_len = payload_length > sizeof(data) ? sizeof(data) : payload_length;
    if (bpf_skb_load_bytes(skb, payload_offset, data, data_len) < 0) {
        return 0;
    }
    
    // 检查是否为HTTP
    int http_type = is_http(data, data_len);
    if (http_type == 0) {
        return 0;
    }
    
    // 如果是新的HTTP会话，添加到会话表
    if (!is_known_session(&session_key)) {
        add_session(&session_key);
    }
    
    // 创建事件
    struct http_event_t event = {0};
    event.ts = bpf_ktime_get_ns();
    event.pid = bpf_get_current_pid_tgid() >> 32;
    event.saddr = iph.saddr;
    event.daddr = iph.daddr;
    event.sport = bpf_ntohs(tcph.source);
    event.dport = bpf_ntohs(tcph.dest);
    event.direction = http_type;
    
    // 复制数据
    int i;
    for (i = 0; i < data_len; i++) {
        event.data[i] = data[i];
    }
    
    // 发送事件
    bpf_perf_event_output(skb, &http_events, BPF_F_CURRENT_CPU, &event, sizeof(event));
    
    return 0;
}
